use ndarray::Array1;
use ndarray::Array2;
use ndarray::s;
use crate::netlist::VoltageKind;

#[derive(Debug, Clone, Copy)]
pub enum NodeIndex {
    GND,
    Node(usize)
}

#[derive(Debug)]
pub struct R(pub NodeIndex, pub NodeIndex, pub f64);

#[derive(Debug)]
pub struct V(pub NodeIndex, pub NodeIndex, pub VValue);

#[derive(Debug)]
pub enum VValue {
    Src(VoltageKind),
    Cap {
        cur_v: f64,
        cap: f64,
    }
}

impl VValue {
    pub fn cur_value(&self, time: f64) -> f64 {
        match self {
            Self::Src(v) => v.voltage(time),
            Self::Cap { cur_v, cap: _ } => *cur_v,
        }
    }

    pub fn update(&mut self, delta_time: f64, i: f64) {
        if let Self::Cap { cur_v, cap } = self {
            *cur_v -= (delta_time / *cap) * i;
        }
    }
}

#[derive(Debug)]
pub struct I(pub NodeIndex, pub NodeIndex, pub IValue);

#[derive(Debug)]
pub enum IValue {
    Src(f64),
    Ind {
        cur_i: f64,
        lnd: f64,
    }
}

impl IValue {
    pub fn cur_value(&self, _time: f64) -> f64 {
        match self {
            Self::Src(c) => *c,
            Self::Ind { cur_i, lnd: _ } => *cur_i,
        }
    }

    pub fn update(&mut self, delta_time: f64, v: f64) {
        if let Self::Ind { cur_i, lnd } = self {
            *cur_i -= (delta_time / *lnd) * v;
        }
    }
}

#[derive(Debug)]
pub struct MNAData {
    pub rs: Vec<R>,
    pub vs: Vec<V>,
    pub is: Vec<I>,

    pub node_size: usize,

    pub time: f64,
}

impl MNAData {
    pub fn update(&mut self, delta_time: f64, result: &Array1<f64>) {
        let i_result = result.slice(s![self.node_size..]);
        for (vdata, i) in self.vs.iter_mut().zip(i_result.into_iter()) {
            vdata.2.update(delta_time, *i);
        }

        for idata in self.is.iter_mut() {
            let v0 = match idata.0 {
                NodeIndex::GND => 0.,
                NodeIndex::Node(index) => result[index],
            };
    
            let v1 = match idata.1 {
                NodeIndex::GND => 0.,
                NodeIndex::Node(index) => result[index],
            };

            idata.2.update(delta_time, v0 - v1);
        }

        self.time += delta_time;
    }

    pub fn reach(&self, sim_time: f64) -> bool {
        self.time >= sim_time
    }
}

impl MNAData {
    pub fn build_mna(&self) -> (Array2<f64>, Array1<f64>) {
        let g = self.build_g();
        let e = self.build_e();
        let net = -1. * &e.t();
        let y = self.build_y();

        // Build Gx
        let size = self.node_size + self.vs.len();
        let mut gx = Array2::<f64>::zeros((size, size));

        gx.slice_mut(s![0..self.node_size, 0..self.node_size]).assign(&g);
        gx.slice_mut(s![0..self.node_size, self.node_size..size]).assign(&e);
        gx.slice_mut(s![self.node_size..size, 0..self.node_size]).assign(&net);
        
        (gx, y)
    }

    fn build_g(&self) -> Array2<f64> {
        let g_size = self.node_size;
        let mut g = Array2::<f64>::zeros((g_size, g_size));
        
        for r in self.rs.iter() {
            let conduction = 1.0 / r.2;
            match (r.0, r.1) {
                (NodeIndex::GND, NodeIndex::GND) => {
                    panic!("both 0")
                }
                (NodeIndex::Node(i1), NodeIndex::Node(i2)) => {
                    g[[i1, i2]] += -conduction; 
                    g[[i2, i1]] += -conduction; 
                    g[[i1, i1]] +=  conduction; 
                    g[[i2, i2]] +=  conduction;        
                }
                (NodeIndex::GND, NodeIndex::Node(index)) => {
                    g[[index, index]] += conduction;
                }
                (NodeIndex::Node(index), NodeIndex::GND) => {
                    g[[index, index]] += conduction;
                }
            }
        }

        g
    }

    fn build_e(&self) -> Array2<f64> {
        let row_size = self.node_size;
        let col_size = self.vs.len();
        let mut e = Array2::<f64>::zeros((row_size, col_size));

        for (i, v) in self.vs.iter().enumerate() {
            if let NodeIndex::Node(index) = v.0 {
                e[[index, i]] = -1.;
            }
            if let NodeIndex::Node(index) = v.1 {
                e[[index, i]] = 1.;
            }
        }

        e
    }

    fn build_y(&self) -> Array1<f64> {
        let row_size = self.node_size + self.vs.len();
        let mut y = Array1::<f64>::zeros(row_size);

        for i in self.is.iter() {
            if let NodeIndex::Node(index) = i.0 {
                y[index] = i.2.cur_value(self.time);
            }
            if let NodeIndex::Node(index) = i.1 {
                y[index] = -i.2.cur_value(self.time);
            }
        }

        for (i, v) in self.vs.iter().enumerate() {
            y[self.node_size + i] = v.2.cur_value(self.time);
        }

        y
    }

}